# -*- coding: utf-8 -*-
"""
临时加密网页聊天（含本地密文存储 + 设备口令离线回看）
- FastAPI + Uvicorn
- 浏览器端 ECDH(P-256) + HKDF + AES-GCM（端到端加密）
- 可选入口口令（PBKDF2-HMAC-SHA256 100k，服务端仅存盐+哈希）
- 链接按天过期；握手重播；消息队列；本地 http 兼容
- IndexedDB 只存密文（IV+DATA）；可用“设备口令”对会话密钥封装保存，离线回看
"""
import asyncio, json, secrets, time
from contextlib import asynccontextmanager
from datetime import datetime, timedelta
from typing import Dict, Optional

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
from fastapi.responses import HTMLResponse, JSONResponse, Response

# ---------------- Room ----------------
class Room:
    def __init__(self, room_id: str, expires_at: float):
        self.room_id = room_id
        self.expires_at = expires_at
        self.owner: Optional[WebSocket] = None
        self.guest: Optional[WebSocket] = None
        self.owner_name = "发布者"
        self.guest_name = "访客"
        self.pass_salt_b64: Optional[str] = None
        self.pass_hash_b64: Optional[str] = None
        # 握手缓存（用于重播）
        self.last_salt: Optional[list] = None      # [int,...]
        self.owner_pubkey: Optional[dict] = None   # JWK
        self.guest_pubkey: Optional[dict] = None   # JWK

    @property
    def expired(self) -> bool:
        return time.time() > self.expires_at

    def other(self, ws: WebSocket) -> Optional[WebSocket]:
        return self.guest if ws is self.owner else (self.owner if ws is self.guest else None)

ROOMS: Dict[str, Room] = {}

async def gc_rooms():
    while True:
        dead = [rid for rid, r in list(ROOMS.items()) if r.expired or (r.owner is None and r.guest is None)]
        for rid in dead:
            ROOMS.pop(rid, None)
        await asyncio.sleep(30)

@asynccontextmanager
async def lifespan(app: FastAPI):
    asyncio.create_task(gc_rooms()); yield

app = FastAPI(lifespan=lifespan)

# ---------------- Pages ----------------
INDEX_HTML = """
<!doctype html><html lang="zh-CN"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>密聊闪信 · 临时加密网页聊天</title>
<meta name="description" content="网页版临时加密聊天：通过浏览器直接使用，无需App，端到端加密，到期自毁。">
<meta name="theme-color" content="#0b1220">
<link rel="icon" href="data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 64 64'%3E%3Ccircle cx='32' cy='32' r='28' fill='%230b1220'/%3E%3Cpath d='M20 30v10h24V30h-2v-4a10 10 0 10-20 0v4h-2zm6-4a6 6 0 1112 0v4H26v-4z' fill='%237cdcff'/%3E%3C/svg%3E">
<style>
body{margin:0;background:#0b1220;color:#e8ecf1;font:16px/1.6 system-ui,Segoe UI,Roboto,Helvetica,Arial}
.wrap{min-height:100vh;display:grid;place-items:center;padding:40px}
.card{width:min(680px,100%);background:#101a2a;border:1px solid rgba(255,255,255,.12);border-radius:16px;padding:28px;box-shadow:0 10px 40px rgba(0,0,0,.4)}
h1{margin:0 0 6px;font-weight:700}.sub{color:#7f8da9;margin-bottom:8px}
.badges{display:flex;flex-wrap:wrap;gap:8px;margin:8px 0 16px}
.badge{font-size:12px;padding:2px 8px;border:1px solid rgba(255,255,255,.2);border-radius:999px;background:rgba(124,220,255,.12)}
label{display:block;margin:14px 0 6px}
input,button{width:100%;padding:14px 16px;border-radius:12px;border:1px solid rgba(255,255,255,.12);background:rgba(255,255,255,.04);color:#e8ecf1}
input:focus{outline:2px solid #7cdcff;border-color:transparent}
button{cursor:pointer;background:linear-gradient(90deg,#2aa3ff,#7cdcff);color:#00111a;font-weight:700;border:none}
.row{display:grid;grid-template-columns:1fr auto;gap:10px}
.link{margin-top:14px;padding:12px;border-radius:12px;background:rgba(255,255,255,.06);word-break:break-all}
.tip{color:#7f8da9;font-size:14px;margin-top:10px}
footer{margin-top:16px;text-align:center;color:#7f8da9;font-size:12px}
</style></head>
<body><div class="wrap"><div class="card">
<h1>密聊闪信 · Web</h1>
<div class="sub">通过浏览器发起一次性密聊链接，对方打开网页即可聊天（无需下载App）。</div>
<div class="badges">
  <span class="badge">网页聊天</span><span class="badge">端到端加密</span>
  <span class="badge">临时链接</span><span class="badge">到期自毁</span>
  <span class="badge">可选口令</span>
</div>
<label>过期天数（1-30）</label>
<div class="row"><input id="days" type="number" min="1" max="30" value="3"><button id="create">生成临时加密网页链接</button></div>
<label>进入口令（可选，不会出现在链接中）</label>
<input id="pass" type="password" placeholder="留空则无需口令">
<div id="out" class="link" style="display:none"></div>
<div class="tip">提示：若设置了口令，请线下告知对方。</div>
<footer>这是网页版密聊工具，无需安装 App，使用浏览器即可。</footer>
</div></div>
<script>
(function(){
  var createBtn=document.getElementById('create'); var out=document.getElementById('out');
  createBtn.onclick=function(){
    var days=Math.max(1,Math.min(30,parseInt(document.getElementById('days').value||'3')));
    // 口令统一做 NFKC + trim，避免全角/空白差异
    var pass=(document.getElementById('pass').value||''); if(pass.normalize) pass=pass.normalize('NFKC'); pass=pass.trim();

    fetch('/api/create',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({days:days})})
      .then(function(r){return r.json();})
      .then(function(j){
        var room=j.room;
        function afterPass(){
          var url=location.origin+'/chat?room='+encodeURIComponent(room)+'&exp='+encodeURIComponent(j.expires_at);
          out.style.display='block'; out.textContent=url;
          if(navigator.clipboard && navigator.clipboard.writeText){ navigator.clipboard.writeText(url).catch(function(){}); }
        }
        if(pass){
          var salt=new Uint8Array(16); window.crypto.getRandomValues(salt);
          var enc=new TextEncoder();
          window.crypto.subtle.importKey('raw',enc.encode(pass),'PBKDF2',false,['deriveBits']).then(function(km){
            return window.crypto.subtle.deriveBits({name:'PBKDF2',hash:'SHA-256',salt:salt,iterations:100000},km,256);
          }).then(function(bits){
            var hash_b64=btoa(String.fromCharCode.apply(null,new Uint8Array(bits)));
            var salt_b64=btoa(String.fromCharCode.apply(null,salt));
            return fetch('/api/set_pass',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({room:room,salt_b64:salt_b64,hash_b64:hash_b64})});
          }).then(function(){ afterPass(); });
        } else { afterPass(); }
      });
  };
})();
</script></body></html>
"""

CHAT_HTML = """
<!doctype html><html lang="zh-CN"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>密聊闪信 · 临时加密网页聊天</title>
<meta name="description" content="网页版临时加密聊天：通过浏览器直接使用，无需App，端到端加密，到期自毁。">
<meta name="theme-color" content="#08111f">
<link rel="icon" href="data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 64 64'%3E%3Ccircle cx='32' cy='32' r='28' fill='%2308111f'/%3E%3Cpath d='M20 30v10h24V30h-2v-4a10 10 0 10-20 0v4h-2zm6-4a6 6 0 1112 0v4H26v-4z' fill='%237cdcff'/%3E%3C/svg%3E">
<style>
body{margin:0;background:#08111f;color:#e8ecf1;font:16px/1.6 system-ui,Segoe UI,Roboto,Helvetica,Arial;height:100vh;display:flex;flex-direction:column}
header{display:flex;align-items:center;gap:10px;padding:14px 16px;border-bottom:1px solid rgba(255,255,255,.1)}
.dot{width:10px;height:10px;border-radius:50%;background:#777}.on{background:#29d17f}
.title{font-weight:700}.meta{color:#8aa0b8;font-size:12px}
main{flex:1;display:flex;flex-direction:column;gap:12px;padding:16px;overflow:auto}
.bubble{max-width:70%;padding:10px 12px;border-radius:14px;margin:4px 0;background:rgba(255,255,255,.06);border:1px solid rgba(255,255,255,.08)}
.me{align-self:flex-end;background:rgba(124,220,255,.15);border-color:rgba(124,220,255,.25)}
.peer{align-self:flex-start}
.composer{display:grid;grid-template-columns:1fr auto auto auto;gap:8px;padding:12px 16px;border-top:1px solid rgba(255,255,255,.08)}
input,button{padding:12px 14px;border-radius:12px;border:1px solid rgba(255,255,255,.12);background:rgba(255,255,255,.04);color:#e8ecf1}
input:focus{outline:2px solid #7cdcff;border-color:transparent}
button{cursor:pointer;background:linear-gradient(90deg,#2aa3ff,#7cdcff);color:#00111a;font-weight:700;border:none}
</style></head>
<body>
<header><div id="dot" class="dot"></div><div class="title">密聊闪信 · Web</div><div style="margin-left:auto" class="meta" id="status">等待连接...</div></header>
<div class="meta" style="padding:6px 16px">若房间设置了口令，请先输入：</div>
<div style="padding:0 16px 8px;display:flex;gap:8px"><input id="passphrase" type="password" placeholder="进入口令（若需要）" style="flex:1"><button id="unlock">进入</button></div>
<main id="log"></main>
<div class="composer">
  <input id="msg" placeholder="输入消息，回车发送">
  <button id="send">发送</button>
  <button id="clear">清屏</button>
  <button id="wipe">清除本地记录</button>
</div>

<script>
(function(){
  function $(id){return document.getElementById(id);}
  function uiAdd(t,cls){var b=document.createElement('div');b.className='bubble '+(cls||'');b.textContent=t;$('log').appendChild(b);b.scrollIntoView({behavior:'smooth',block:'end'});}
  function setPeerOnline(on){$('dot').className='dot '+(on?'on':'');$('status').textContent=on?'已连接':'等待连接...';}

  // ---- 参数
  var href=location.href; var a=href.split('?')[1]||''; var q={},p=a.split('&'); for(var i=0;i<p.length;i++){ if(!p[i])continue; var kv=p[i].split('='); q[decodeURIComponent(kv[0])]=decodeURIComponent(kv[1]||''); }
  var room=q.room||''; if(!room){ alert('缺少 room'); location.href='/'; return; }

  // ===== IndexedDB（仅存密文 + 封装后的会话密钥）=====
  var DB_NAME='mlsx', DB_VER=2, dbPromise=null;
  function idbOpen(){ if(dbPromise) return dbPromise; dbPromise=new Promise(function(resolve,reject){ var req=indexedDB.open(DB_NAME,DB_VER);
    req.onupgradeneeded=function(e){ var db=e.target.result;
      if(!db.objectStoreNames.contains('messages')){ var st=db.createObjectStore('messages',{keyPath:'id'}); st.createIndex('room_ts',['room','ts']); st.createIndex('room','room'); }
      if(!db.objectStoreNames.contains('rooms')){ db.createObjectStore('rooms',{keyPath:'room'}); }
    };
    req.onsuccess=function(e){ resolve(e.target.result); }; req.onerror=function(e){ reject(e.target.error); };
  }); return dbPromise; }
  function idbPutMessage(room,rec){ return idbOpen().then(function(db){ return new Promise(function(resolve,reject){ var tx=db.transaction('messages','readwrite'); tx.oncomplete=function(){resolve();}; tx.onerror=function(e){reject(e.target.error);}; tx.objectStore('messages').put(rec); }); }); }
  function idbLoadHistory(room,limit){ return idbOpen().then(function(db){ return new Promise(function(resolve,reject){ var tx=db.transaction('messages','readonly'); var idx=tx.objectStore('messages').index('room_ts'); var range=IDBKeyRange.bound([room,0],[room,9e15]); var res=[], req=idx.openCursor(range); req.onsuccess=function(e){ var c=e.target.result; if(c){ res.push(c.value); c.continue(); } else { if(typeof limit==='number') res=res.slice(Math.max(0,res.length-limit)); resolve(res); } }; req.onerror=function(e){ reject(e.target.error); }; }); }); }
  function idbClearRoom(room){ return idbOpen().then(function(db){ return new Promise(function(resolve,reject){ var tx=db.transaction('messages','readwrite'); var st=tx.objectStore('messages'); var idx=st.index('room'); var req=idx.openCursor(IDBKeyRange.only(room)); req.onsuccess=function(e){ var c=e.target.result; if(c){ st.delete(c.primaryKey); c.continue(); } else resolve(); }; req.onerror=function(e){ reject(e.target.error); }; }); }); }
  function idbSaveWrappedKey(room, wrapped){ return idbOpen().then(function(db){ return new Promise(function(resolve,reject){ var tx=db.transaction('rooms','readwrite'); tx.oncomplete=function(){resolve();}; tx.onerror=function(e){reject(e.target.error);}; tx.objectStore('rooms').put({room:room, wrappedKey:wrapped}); }); }); }
  function idbLoadWrappedKey(room){ return idbOpen().then(function(db){ return new Promise(function(resolve,reject){ var tx=db.transaction('rooms','readonly'); var req=tx.objectStore('rooms').get(room); req.onsuccess=function(e){ var r=e.target.result; resolve(r? r.wrappedKey : null); }; req.onerror=function(e){ reject(e.target.error); }; }); }); }
  function makeId(){ return Date.now().toString(36)+'-'+Math.random().toString(36).slice(2,8); }

  // ---- 派生设备口令 KEK、封装/解封装房间密钥
  var enc=new TextEncoder(), dec=new TextDecoder();
  function deriveKEKFromDevicePass(pass){
    if(pass.normalize) pass=pass.normalize('NFKC'); pass=pass.trim();
    var salt = enc.encode('mlsx-device-kek-v1'); // 固定盐用于区分用途
    return window.crypto.subtle.importKey('raw', enc.encode(pass), 'PBKDF2', false, ['deriveKey']).then(function(base){
      return window.crypto.subtle.deriveKey(
        {name:'PBKDF2', hash:'SHA-256', salt:salt, iterations:200000},
        base,
        {name:'AES-GCM', length:256},
        false,
        ['encrypt','decrypt']
      );
    });
  }
  function wrapRoomKeyWithKEK(aesKey, devicePass){
    return deriveKEKFromDevicePass(devicePass).then(function(kek){
      return window.crypto.subtle.exportKey('raw', aesKey).then(function(raw){
        var iv=new Uint8Array(12); crypto.getRandomValues(iv);
        return window.crypto.subtle.encrypt({name:'AES-GCM',iv:iv},kek,new Uint8Array(raw)).then(function(ct){
          return {iv:Array.prototype.slice.call(iv), data:Array.prototype.slice.call(new Uint8Array(ct))};
        });
      });
    });
  }
  function unwrapRoomKeyWithKEK(wrapped, devicePass){
    return deriveKEKFromDevicePass(devicePass).then(function(kek){
      var iv=new Uint8Array(wrapped.iv), data=new Uint8Array(wrapped.data);
      return window.crypto.subtle.decrypt({name:'AES-GCM',iv:iv},kek,data).then(function(raw){
        return window.crypto.subtle.importKey('raw', new Uint8Array(raw), 'AES-GCM', true, ['encrypt','decrypt']);
      });
    });
  }

  // ---- E2EE 会话状态
  var myECDHPriv=null, myECDHPubJwk=null, aesKey=null, gotPeerKey=null, gotSalt=null, peerName='对方';
  var myName=localStorage.getItem('nickname')||''; var myRole=sessionStorage.getItem('role:'+room);
  if(!myRole){ myRole = window.confirm('你是发布者吗? (OK=发布者, Cancel=访客)')?'owner':'guest'; sessionStorage.setItem('role:'+room,myRole); }
  if(!myName){ myName=window.prompt('输入你的昵称（可留空）')||''; localStorage.setItem('nickname',myName); }
  var pendingCipher=[], pendingOut=[];
  var needPass=false, lastSaltB64=null;

  // 页面加载：先把本地密文历史塞入待解密队列
  idbLoadHistory(room, 300).then(function(items){
    for(var i=0;i<items.length;i++){
      var dir = items[i].dir || 'in';
      pendingCipher.push({ payload: items[i].payload, from: (dir==='out'?'me':'peer'), ts: items[i].ts, _isHistory: true });
    }
  }).catch(function(){});

  // ECDH & 会话密钥导出允许
  function genECDH(){
    return window.crypto.subtle.generateKey({name:'ECDH',namedCurve:'P-256'},true,['deriveBits']).then(function(kp){
      myECDHPriv=kp.privateKey; return window.crypto.subtle.exportKey('jwk',kp.publicKey);
    }).then(function(pub){ myECDHPubJwk=pub; });
  }
  function deriveAES(peerPubJwk,saltBuf){
    return window.crypto.subtle.importKey('jwk',peerPubJwk,{name:'ECDH',namedCurve:'P-256'},false,[]).then(function(peerPub){
      return window.crypto.subtle.deriveBits({name:'ECDH',public:peerPub},myECDHPriv,256);
    }).then(function(bits){
      return window.crypto.subtle.importKey('raw',bits,'HKDF',false,['deriveKey']);
    }).then(function(ikm){
      // extractable:true 便于本机封装保存
      return window.crypto.subtle.deriveKey(
        {name:'HKDF',hash:'SHA-256',salt:saltBuf,info:enc.encode('temp-chat-v1')},
        ikm,
        {name:'AES-GCM',length:256},
        true,
        ['encrypt','decrypt']
      );
    }).then(function(key){ aesKey=key; onReady(); });
  }
  function encryptMessage(plain){
    var iv=new Uint8Array(12); window.crypto.getRandomValues(iv);
    return window.crypto.subtle.encrypt({name:'AES-GCM',iv:iv},aesKey,enc.encode(plain)).then(function(ct){
      return {iv:Array.prototype.slice.call(iv),data:Array.prototype.slice.call(new Uint8Array(ct))};
    });
  }
  function decryptMessage(p){
    var iv=new Uint8Array(p.iv); var data=new Uint8Array(p.data);
    return window.crypto.subtle.decrypt({name:'AES-GCM',iv:iv},aesKey,data).then(function(pt){ return dec.decode(pt); });
  }

  function onReady(){
    $('status').textContent='已连接（端到端加密）'; setPeerOnline(true);

    // 首次 ready：如果未保存封装密钥，询问是否启用离线回看
    if (!sessionStorage.getItem('offline-enabled:'+room)) {
      idbLoadWrappedKey(room).then(function(w){
        if (w) return;
        setTimeout(function(){
          var ans = confirm('启用本机“离线回看”吗？（会要求设置一个“设备口令”，用于加密保存会话密钥）');
          if (!ans) return;
          var pass = prompt('设置设备口令（仅本机，忘记将无法离线查看）：');
          if (!pass) return;
          wrapRoomKeyWithKEK(aesKey, pass).then(function(wrapped){
            return idbSaveWrappedKey(room, wrapped);
          }).then(function(){
            sessionStorage.setItem('offline-enabled:'+room,'1');
            console.log('[offline] room key wrapped & saved');
          }).catch(function(e){ console.warn('wrap/save failed',e); });
        }, 300);
      });
    }

    // flush outgoing
    while(pendingOut.length){
      (function(v){ encryptMessage(v).then(function(payload){ ws.send(JSON.stringify({type:'cipher',payload:payload})); }); })(pendingOut.shift());
    }
    // flush incoming (历史 + 实时)
    (function drain(){
      if(!pendingCipher.length) return;
      var item=pendingCipher.shift();
      var p=item.payload||item;
      decryptMessage(p).then(function(text){
        var who=(item.from==='me')?'我':(peerName||'对方');
        var cls=(item.from==='me')?'me':'peer';
        var prefix=item._isHistory?'[历史] ':'';
        uiAdd(prefix+who+': '+text,cls);
      }).finally(drain);
    })();
  }

  function proceedHandshake(){
    if(ws && ws.readyState===WebSocket.OPEN){
      ws.send(JSON.stringify({type:'hello',as:myRole,name:myName}));
      if(myRole==='owner'){
        var salt=new Uint8Array(16); window.crypto.getRandomValues(salt);
        ws.send(JSON.stringify({type:'salt',data:Array.prototype.slice.call(salt)}));
        sessionStorage.setItem('salt:'+room, JSON.stringify(Array.prototype.slice.call(salt)));
      }
      ws.send(JSON.stringify({type:'pubkey',data:myECDHPubJwk}));
    }
  }

  if(!window.crypto || !window.crypto.subtle){ alert('当前浏览器不支持 WebCrypto'); return; }
  var ws=null;
  genECDH().then(function(){
    var scheme=(location.protocol==='https:'?'wss':'ws');
    ws=new WebSocket(scheme+'://'+location.host+'/ws/'+room);
    ws.onmessage=function(ev){
      var msg=JSON.parse(ev.data);

      if(msg.type==='need_pass'){
        needPass=true; lastSaltB64=msg.salt_b64; $('status').textContent='需要口令';
        var passNow=$('passphrase').value||''; if(passNow.normalize) passNow=passNow.normalize('NFKC'); passNow=passNow.trim();
        if(passNow){ submitPass(passNow); }
        return;
      }
      if(msg.type==='pass_ok'){ needPass=false; return; }
      if(msg.type==='pass_bad'){ $('status').textContent='口令错误，请重试'; return; }
      if(msg.type==='ready'){ proceedHandshake(); return; }
      if(msg.type==='error'){ alert('错误: '+msg.reason); location.href='/'; return; }
      if(msg.type==='peer_status'){ setPeerOnline(!!msg.online); if(msg.name) peerName=msg.name; return; }

      if(msg.type==='salt'){
        gotSalt=new Uint8Array(msg.data);
        sessionStorage.setItem('salt:'+room, JSON.stringify(msg.data));
        if(gotPeerKey&&gotSalt){ deriveAES(gotPeerKey,gotSalt); }
        return;
      }
      if(msg.type==='pubkey'){
        gotPeerKey=msg.data;
        var s=sessionStorage.getItem('salt:'+room); if(s&&!gotSalt){ gotSalt=new Uint8Array(JSON.parse(s)); }
        if(gotPeerKey&&gotSalt){ deriveAES(gotPeerKey,gotSalt); }
        return;
      }
      if(msg.type==='cipher'){
        // 收到密文：先写入本地（加密态），再显示/或待解密
        var rec={id:makeId(),room:room,ts:Date.now(),dir:'in',payload:msg.payload};
        idbPutMessage(room,rec).catch(function(){});
        if(!aesKey){ pendingCipher.push({payload:msg.payload,from:'peer',ts:rec.ts}); return; }
        decryptMessage(msg.payload).then(function(text){ uiAdd((peerName||'对方')+': '+text,'peer'); });
        return;
      }
    };

    // 如果 2 秒内还没有 ready（比如对方不在线），尝试离线恢复密钥以查看本地历史
    setTimeout(function(){
      if (aesKey) return;
      idbLoadWrappedKey(room).then(function(w){
        if (!w) return;
        var pass = prompt('输入“设备口令”以离线查看本地历史：');
        if (!pass) return;
        unwrapRoomKeyWithKEK(w, pass).then(function(key){
          aesKey = key;
          onReady(); // 用离线密钥解密 pendingCipher（包含本地历史）
        }).catch(function(){ alert('设备口令不正确，无法离线查看。'); });
      });
    }, 2000);
  });

  // 主动提交入口口令
  function submitPass(pass){
    if(!needPass || !lastSaltB64) return;
    if(pass.normalize) pass=pass.normalize('NFKC'); pass=pass.trim();
    if(!pass){ $('status').textContent='需要口令'; return; }
    var raw=atob(lastSaltB64); var salt=new Uint8Array(raw.length); for(var i=0;i<raw.length;i++){ salt[i]=raw.charCodeAt(i); }
    window.crypto.subtle.importKey('raw',enc.encode(pass),'PBKDF2',false,['deriveBits']).then(function(km){
      return window.crypto.subtle.deriveBits({name:'PBKDF2',hash:'SHA-256',salt:salt,iterations:100000},km,256);
    }).then(function(bits){
      var hash_b64=btoa(String.fromCharCode.apply(null,new Uint8Array(bits)));
      ws.send(JSON.stringify({type:'pass_proof',hash_b64:hash_b64}));
      $('status').textContent='正在验证口令…';
    });
  }

  $('unlock').onclick=function(){
    var pass=$('passphrase').value||'';
    if(needPass){ submitPass(pass); return; }
    proceedHandshake();
  };

  var input=$('msg'), sendBtn=$('send'), clearBtn=$('clear'), wipeBtn=$('wipe');
  function send(){
    var v=(input.value||'').trim(); if(!v) return;
    if(!aesKey){ pendingOut.push(v); uiAdd('（待加密发送）: '+v,'me'); input.value=''; return; }
    encryptMessage(v).then(function(payload){
      ws.send(JSON.stringify({type:'cipher',payload:payload}));
      // 本地保存密文记录（不落明文）
      var rec={id:makeId(),room:room,ts:Date.now(),dir:'out',payload:payload};
      idbPutMessage(room,rec).catch(function(){});
      uiAdd('我: '+v,'me'); input.value='';
    });
  }
  sendBtn.onclick=send; input.addEventListener('keydown',function(e){ if(e.key==='Enter'){ send(); }});
  clearBtn.onclick=function(){ $('log').innerHTML=''; };
  wipeBtn.onclick=function(){
    if(!confirm('仅清除此浏览器本地的聊天记录（密文），无法恢复。确定吗？')) return;
    idbClearRoom(room).then(function(){ $('log').innerHTML=''; uiAdd('已清除本地记录（不影响对方与服务器密文）','peer'); });
  };
})();
</script>
</body></html>
"""

# ---------------- Routes ----------------
@app.get("/")
async def index():
    return HTMLResponse(INDEX_HTML)

@app.get("/chat")
async def chat():
    return HTMLResponse(CHAT_HTML)

# 可选：消音 DevTools 探测日志
@app.get("/.well-known/appspecific/com.chrome.devtools.json")
async def chrome_probe():
    return Response(status_code=204)

@app.post("/api/create")
async def api_create(request: Request):
    data = await request.json()
    days = max(1, min(int(data.get("days", 1)), 30))
    rid = secrets.token_urlsafe(10)
    expires_at = (datetime.utcnow() + timedelta(days=days)).timestamp()
    ROOMS[rid] = Room(rid, expires_at)
    return JSONResponse({"room": rid, "expires_at": int(expires_at)})

@app.post("/api/set_pass")
async def api_set_pass(request: Request):
    data = await request.json()
    rid = data.get("room")
    room = ROOMS.get(rid)
    if not room or room.expired:
        return JSONResponse({"ok": False, "reason": "expired_or_not_found"}, status_code=400)
    room.pass_salt_b64 = data.get("salt_b64")
    room.pass_hash_b64 = data.get("hash_b64")
    return JSONResponse({"ok": True})

# ---------------- WebSocket ----------------
@app.websocket("/ws/{room_id}")
async def ws_endpoint(websocket: WebSocket, room_id: str):
    await websocket.accept()
    room = ROOMS.get(room_id)
    if not room or room.expired:
        await websocket.send_text(json.dumps({"type":"error","reason":"expired_or_not_found"})); await websocket.close(); return

    # 口令挑战（允许重试）
    if room.pass_hash_b64 and room.pass_salt_b64:
        await websocket.send_text(json.dumps({"type":"need_pass","salt_b64":room.pass_salt_b64}))
        while True:
            try:
                proof_msg = json.loads(await websocket.receive_text())
            except Exception:
                await websocket.close(); return
            if proof_msg.get("type")=="pass_proof":
                if proof_msg.get("hash_b64")==room.pass_hash_b64:
                    await websocket.send_text(json.dumps({"type":"pass_ok"}))
                    await websocket.send_text(json.dumps({"type":"ready"}))
                    break
                else:
                    await websocket.send_text(json.dumps({"type":"pass_bad"}))
            # 其它消息忽略，直到 pass 通过
    else:
        await websocket.send_text(json.dumps({"type":"ready"}))

    # Hello & 进房
    who = None
    try:
        hello = json.loads(await websocket.receive_text())
        if hello.get("type")!="hello" or hello.get("as") not in ("owner","guest"):
            await websocket.send_text(json.dumps({"type":"error","reason":"bad_handshake"})); await websocket.close(); return
        who = hello["as"]; name = hello.get("name") or ("发布者" if who=="owner" else "访客")
        if who=="owner":
            if room.owner is not None:
                await websocket.send_text(json.dumps({"type":"error","reason":"owner_exists"})); await websocket.close(); return
            room.owner = websocket; room.owner_name = name
        else:
            if room.guest is not None:
                await websocket.send_text(json.dumps({"type":"error","reason":"room_full"})); await websocket.close(); return
            room.guest = websocket; room.guest_name = name

        # 在线通知
        other = room.other(websocket)
        if other:
            await other.send_text(json.dumps({"type":"peer_status","online":True,"name":name}))
            await websocket.send_text(json.dumps({"type":"peer_status","online":True,"name": room.owner_name if who=="guest" else room.guest_name}))
        else:
            await websocket.send_text(json.dumps({"type":"peer_status","online":False}))

        # 重播握手材料（确保两边都拿到 salt 与对方 pubkey）
        async def replay(room_obj: Room):
            if room_obj.owner and room_obj.guest:
                if room_obj.last_salt:
                    await room_obj.owner.send_text(json.dumps({"type":"salt","data": room_obj.last_salt}))
                    await room_obj.guest.send_text(json.dumps({"type":"salt","data": room_obj.last_salt}))
                if room_obj.owner_pubkey:
                    await room_obj.guest.send_text(json.dumps({"type":"pubkey","data": room_obj.owner_pubkey}))
                if room_obj.guest_pubkey:
                    await room_obj.owner.send_text(json.dumps({"type":"pubkey","data": room_obj.guest_pubkey}))

        await replay(room)

        # 主循环：缓存 & 转发
        while True:
            payload = json.loads(await websocket.receive_text())
            t = payload.get("type")
            if t in {"pubkey","salt","cipher","typing","bye"}:
                if t == "salt":
                    room.last_salt = payload.get("data")
                elif t == "pubkey":
                    if websocket is room.owner: room.owner_pubkey = payload.get("data")
                    elif websocket is room.guest: room.guest_pubkey = payload.get("data")
                peer = room.other(websocket)
                if peer: await peer.send_text(json.dumps(payload))
                else:    await replay(room)

    except WebSocketDisconnect:
        pass
    finally:
        if room := ROOMS.get(room_id):
            if room.owner is websocket: room.owner=None
            if room.guest is websocket: room.guest=None
            other = room.other(websocket)
            if other:
                try: await other.send_text(json.dumps({"type":"peer_status","online":False}))
                except Exception: pass

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
